package net.cyue.ort.llm.util;

import net.cyue.ort.llm.generator.GenerationConfig;
import net.cyue.ort.llm.tokenizer.Tokenizer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Arrays;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 默认Token管理器实现
 * 封装现有的Tokenizer接口
 */
public class DefaultTokenManager implements TokenManager {
    private static final Logger logger = LoggerFactory.getLogger(DefaultTokenManager.class);

    private final Tokenizer tokenizer;
    private long[] eosTokenIDs;
    
    public DefaultTokenManager(Tokenizer tokenizer) {
        this.tokenizer = tokenizer;
    }
    
    public DefaultTokenManager(Tokenizer tokenizer, long[] eosTokenIDs) {
        this.tokenizer = tokenizer;
        this.eosTokenIDs = eosTokenIDs;
    }
    
    @Override
    public List<Long> tokenize(String text) {
        long startTime = System.currentTimeMillis();
        Tokenizer.Result result = tokenizer.encode(text);
        List<Long> ids = Arrays.stream(result.ids).boxed().collect(Collectors.toList());
        logTokenizationResult(text, result, System.currentTimeMillis() - startTime);
        return ids;
    }
    
    @Override
    public String detokenize(List<Long> tokenIds) {
        long startTime = System.currentTimeMillis();
        String text = tokenizer.decode(tokenIds);
        if (logger.isDebugEnabled()) {
            logger.debug(
                "Detokenize completed: tokens={}, outputLength={}, time={}ms",
                tokenIds == null ? 0 : tokenIds.size(),
                text == null ? 0 : text.length(),
                System.currentTimeMillis() - startTime
            );
        }
        return text;
    }
    
    @Override
    public boolean isEndOfSequence(long tokenId) {
        if (eosTokenIDs == null) {
            return false;
        }
        boolean matched = Arrays.stream(eosTokenIDs).anyMatch(id -> id == tokenId);
        if (matched && logger.isDebugEnabled()) {
            logger.debug("Detected EOS token ID: {}", tokenId);
        }
        return matched;
    }
    
    @Override
    public boolean isEndOfSequence(List<Long> tokenIds) {
        if (eosTokenIDs == null || tokenIds == null) {
            return false;
        }
        boolean contains = tokenIds.stream().anyMatch(id -> isEndOfSequence(id));
        if (contains && logger.isDebugEnabled()) {
            logger.debug("Token list contains EOS. Token IDs (sample): {}", summarizeTokenIds(tokenIds));
        }
        return contains;
    }
    
    /**
     * 根据生成配置更新EOS token IDs
     */
    public void updateEosTokenIDs(GenerationConfig config) {
        if (config != null && config.getEosTokenIDs() != null) {
            this.eosTokenIDs = config.getEosTokenIDs();
            if (logger.isDebugEnabled()) {
                logger.debug("Updated EOS tokens from GenerationConfig: {}", summarizeIds(eosTokenIDs));
            }
        }
    }
    
    /**
     * 获取底层tokenizer
     */
    public Tokenizer getTokenizer() {
        return tokenizer;
    }

    private String summarizeIds(long[] ids) {
        if (ids == null || ids.length == 0) {
            return "[]";
        }
        int limit = Math.min(ids.length, 5);
        String summary = Arrays.stream(ids, 0, limit)
            .boxed()
            .map(String::valueOf)
            .collect(Collectors.joining(", "));
        if (ids.length > limit) {
            summary += ", ... (total " + ids.length + ")";
        }
        return "[" + summary + "]";
    }

    private String summarizeTokenIds(List<Long> ids) {
        if (ids == null || ids.isEmpty()) {
            return "[]";
        }
        int limit = Math.min(ids.size(), 5);
        String summary = ids.subList(0, limit)
            .stream()
            .map(String::valueOf)
            .collect(Collectors.joining(", "));
        if (ids.size() > limit) {
            summary += ", ... (total " + ids.size() + ")";
        }
        return "[" + summary + "]";
    }

    private void logTokenizationResult(String text, Tokenizer.Result result, long durationMs) {
        if (!logger.isDebugEnabled() || result == null) {
            return;
        }
        String previewTokens = previewArray(result.tokens, 6);
        String previewIds = previewArray(result.ids, 8);
        logger.debug("""

╔════════ Tokenizer Result ════════
║ Input length  : {}
║ Tokens        : {} {}
║ IDs           : {} {}
║ Prompt        : {}
║ Elapsed       : {} ms
╚═════════════════════════════════
""",
            text == null ? 0 : text.length(),
            result.tokens == null ? 0 : result.tokens.length,
            previewTokens,
            result.ids == null ? 0 : result.ids.length,
            previewIds,
            result.prompt == null ? "(none)" : abbreviate(result.prompt, 80),
            durationMs
        );
    }

    private String previewArray(String[] items, int limit) {
        if (items == null || items.length == 0) {
            return "[]";
        }
        String content = Arrays.stream(items, 0, Math.min(items.length, limit))
            .map(s -> s == null ? "null" : abbreviate(s, 20))
            .collect(Collectors.joining(", "));
        if (items.length > limit) {
            content += ", … (" + items.length + ")";
        }
        return "[" + content + "]";
    }

    private String previewArray(long[] items, int limit) {
        if (items == null || items.length == 0) {
            return "[]";
        }
        StringBuilder sb = new StringBuilder("[");
        for (int i = 0; i < Math.min(items.length, limit); i++) {
            if (i > 0) sb.append(", ");
            sb.append(items[i]);
        }
        if (items.length > limit) {
            sb.append(", … (").append(items.length).append(")");
        }
        sb.append("]");
        return sb.toString();
    }

    private String abbreviate(String text, int maxLen) {
        if (text == null || text.length() <= maxLen) {
            return text;
        }
        return text.substring(0, Math.max(0, maxLen - 3)) + "...";
    }
}

